import datetime
import numpy as np
import Core.Gadget as Gadget
import pandas as pd
import Core.MySQLDB as MySQLDB


class THSDatabase(MySQLDB.MySQLDB):
    #
    def __init__(self, address, port, username, password, dbname="sam_ths"):
        # super().__init__(address, port, username=username, password=password)

        self.address = address
        self.port = port
        self.username = username
        self.password = password
        self.contexts = {}
        #
        self.dbname = dbname  # "sam_ths"
        self.name = "THSDatabase THS数据库"
        print("建立数据库对象", self.name)

        # sec_basic_info 中的字段是 sec_issuer_id
        # 公募基金
        # sec_issuer_id -> fund_code (KF0003)，
        # 基础信息表 fund_code 索引，
        # 日数据用 fund_code 索引

        # 指数
        # sec_issuer_id -> index_id (ind249), sec_code->index_code ，
        # 指数基础数据 index_code(000001) 和 index_id 索引，
        # 指数日数据用 index_code 和 sec_id(SCSIH11001) 索引

        # 股票对应 sec_issuer_id -> ?? (gssz0000001),
        # 基础信息表 ??
        # 日数据用 sec_code(000001) 和 sec_id(SSZ000001) 索引

        #
        self.fund_code_by_symbol = {}  # MutualFund: fund_code is index of fund_net_value
        #
        # Index: sec_id is index of index_daily_quotation
        # Stock : sec_code, sec_id is index of shsz_stock_daily_quotation
        self.sec_id_by_symbol = {}

        #
        self.refresh_instrument_map_table_datetime = None
        self._Refresh_Instrument_Map_Table()

    def _Special_Symbol(self):
        config_by_symbol = {}
        # fixed_rate_1y
        # select adjust_date, fixed_rate_1y from  rmb_deposit_benchmark_interest
        # 000007.CZC
        # 000008.CZC

    def _Get_Instrument_Filter(self, instrument_type):
        if instrument_type:
            instrument_type = instrument_type.lower()
        if instrument_type == "mutualfund":
            pass
        elif instrument_type == "stock":
            pass

        elif instrument_type == "index":
            pass

        elif instrument_type == "future":
            pass

        elif instrument_type == "bond":
            pass

        else:
            pass

    # ---建立码表映射关系---
    def _Build_MutualFund_InstrumentTable_MutualFund(self):
        # 基金
        filter = [("sec_type", "开放式基金")]
        filter.append(("or", "sec_type", "封闭式基金"))
        filter.append(("or", "sec_type", "ETF"))
        filter.append(("or", "sec_type", "LOF"))
        filter.append(("or", "sec_type", "QDII"))

        projection = ["sec_issuer_id", "sec_code", "thscode", "sec_short_name_cn", "sec_full_name", "sec_type"]
        instruments = self.Find(self.dbname, "sec_basic_info", filter=filter, projection=projection)
        #
        for instrument in instruments:
            symbol = instrument["thscode"]
            if symbol == None or symbol == "":
                # print("symbol is None", "ths_id", ths_id)
                continue
            ths_id = instrument["sec_issuer_id"]
            #
            if symbol in self.fund_code_by_symbol:
                print("重复 symbol", "symbol", symbol, "ths_id", ths_id)
            else:
                self.fund_code_by_symbol[symbol] = ths_id

    def _Build_Stock_Instrument_Map_Table(self):
        filter = []
        filter.append(("sec_type", "A股"))
        projection = ["sec_issuer_id", "sec_code", "thscode", "sec_id", "sec_short_name_cn", "sec_full_name", "sec_type"]
        instruments = self.Find(self.dbname, "sec_basic_info", filter=filter, projection=projection)
        #
        for instrument in instruments:
            symbol = instrument["thscode"]
            ths_id = instrument["sec_id"]
            if symbol in self.sec_id_by_symbol:
                print("重复symbol", "symbol", symbol, "ths_id", ths_id)
            else:
                self.sec_id_by_symbol[symbol] = ths_id

    def _Build_Index_Instrument_Map_Table(self):
        filter = []
        filter.append(("sec_type", "股票指数"))
        filter.append(("or", "sec_type", "基金指数"))
        filter.append(("or", "sec_type", "债券指数"))
        filter.append(("or", "sec_type", "期货指数"))
        projection = ["sec_issuer_id", "sec_code", "thscode", "sec_id", "sec_short_name_cn", "sec_full_name", "sec_type"]
        instruments = self.Find(self.dbname, "sec_basic_info", filter=filter, projection=projection)
        #
        for instrument in instruments:
            symbol = instrument["thscode"]
            ths_id = instrument["sec_id"]
            if symbol in self.sec_id_by_symbol:
                print("重复symbol", "symbol", symbol, "ths_id", ths_id)
            else:
                self.sec_id_by_symbol[symbol] = ths_id

    def _Build_Bond_Instrument_Map_Table(self):
        # 债券
        bond_code = self.Find(self.dbname, "sec_classi_public_code_table", filter=[("class_encode", "002")], projection=["ctgry_code"])
        bond_code_res = [item[key] for item in bond_code for key in item]
        filter = list(zip(['or'] * len(bond_code_res), ['sec_type_code'] * len(bond_code_res), bond_code_res))
        filter[0] = filter[0][1:]

        projection = ["sec_issuer_id", "sec_code", "thscode", "sec_id", "sec_short_name_cn", "sec_full_name", "sec_type"]
        instruments = self.Find(self.dbname, "sec_basic_info", filter=filter, projection=projection)

        for instrument in instruments:
            symbol = instrument["thscode"]
            ths_id = instrument["sec_id"]
            if symbol in self.sec_id_by_symbol:
                print("重复symbol", "symbol", symbol, "ths_id", ths_id)
            else:
                self.sec_id_by_symbol[symbol] = ths_id

    def _Build_Instrument_Map_Table(self):
        # 基金
        self._Build_MutualFund_InstrumentTable_MutualFund()

        # 股票
        self._Build_Stock_Instrument_Map_Table()

        # 指数
        self._Build_Index_Instrument_Map_Table()
        self.sec_id_by_symbol["AU.SHF"] = "S001463101"  # 南华黄金指数

        # 期货

        # 债券
        self._Build_Bond_Instrument_Map_Table()

        # 必要的补充：

    def _Build_Instrument_Map_Table_Common(self, sec_type_filter={}):
        #
        projection = ["sec_issuer_id", "sec_code", "thscode", "sec_short_name_cn", "sec_full_name", "sec_type"]
        instruments = self.Find(self.dbname, "sec_basic_info", filter=sec_type_filter, projection=projection)
        instruments_no_duplicated = []
        for instrument in instruments:
            symbol = instrument["thscode"]
            ths_id = instrument["sec_issuer_id"]
            #
            if symbol == "" or symbol == None:
                # print("空symbol", symbol, thscode)
                continue
            #
            if symbol in self.thsid_by_symbol:
                print("重复symbol", "symbol", symbol, "thscode", ths_id)
            else:
                self.thsid_by_symbol[symbol] = ths_id

            if ths_id in self.symbol_by_thsid:
                print("重复thscode", "symbol", symbol, "thscode", ths_id)
            else:
                self.symbol_by_thsid[ths_id] = symbol
        #
        self.df_instruments = Gadget.DocumentsToDataFrame(instruments)
        self.df_instruments.rename(columns={"thscode": "symbol"}, inplace=True)
        self.df_instruments.set_index("symbol", inplace=True)

    # 第一次进入，或者更细时间已经超过一天，更新代码表
    def _Refresh_Instrument_Map_Table(self):
        now_datetime = datetime.datetime.now()
        refresh = False
        if self.refresh_instrument_map_table_datetime == None:
            refresh = True
        elif str(now_datetime.date()) != str(self.refresh_instrument_map_table_datetime.date()):
            refresh = True
        #
        if refresh:
            self._Build_Instrument_Map_Table()
            self.refresh_instrument_map_table_datetime = now_datetime

    #
    def Get_Instruments(self, instrument_type, filter=[], projection=[]):
        df_instruments = self.Get_Instruments_DataFrame(instrument_type, filter=filter, projection=projection)
        documents = df_instruments.to_dict(orient='records')
        return documents

    def Get_Instruments_DataFrame(self, instrument_type, filter=[], projection=[], listed=True):
        #
        df_instrument = pd.DataFrame()
        # 转换小写
        if instrument_type:
            instrument_type = instrument_type.lower()
        #
        if instrument_type == "mutualfund":
            df_instrument = self._Get_Instruments_DataFrame_MutualFund(filter=filter, projection=projection, listed=True)
        elif instrument_type == "stock":
            df_instrument = self._Get_Instruments_DataFrame_Stock(filter=filter, projection=projection, listed=True)
        elif instrument_type == "index":
            df_instrument = self._Get_Instruments_DataFrame_Index(filter=filter, projection=projection, listed=True)
        elif instrument_type == "bond":
            df_instrument = self._Get_Instruments_DataFrame_Bond(filter=filter, projection=projection, listed=True)
        elif instrument_type == "future":
            pass

        #
        return df_instrument

    def _Get_Instruments_DataFrame_MutualFund(self, filter=[], projection=[], listed=True):
        basic_filter = [("sec_type", "开放式基金")]
        basic_filter.append(("or", "sec_type", "封闭式基金"))
        basic_filter.append(("or", "sec_type", "ETF"))
        basic_filter.append(("or", "sec_type", "LOF"))
        basic_filter.append(("or", "sec_type", "QDII"))
        #
        documents_symbol = self.Find(self.dbname, "sec_basic_info",
                                     filter=basic_filter,
                                     projection=["sec_issuer_id", "sec_code", "thscode", "sec_id"])
        df_symbol = Gadget.DocumentsToDataFrame(documents_symbol)
        # fund_invest_type
        # 0 股票型
        # 1 债券型
        # 2 混合型
        # 3 货币型
        # 6 基金型
        # filter = [("fund_invest_type", "0"), ("or", "fund_invest_type", "2")]
        documents_instrument = self.Find(self.dbname, "fund_basic_info", filter=filter,
                                         projection=["fund_code", "fund_short_name", "fund_name", "fund_invest_type", "established_date"])
        df_instrument = Gadget.DocumentsToDataFrame(documents_instrument)
        #
        df_instrument = pd.merge(df_instrument, df_symbol, how="left", left_on="fund_code", right_on="sec_issuer_id")
        df_instrument["symbol"] = df_instrument["thscode"]
        return df_instrument

    def _Get_Instruments_DataFrame_Index(self, filter=[], projection=[], listed=True):
        basic_filter = []
        basic_filter.append(("sec_type", "股票指数"))
        basic_filter.append(("or", "sec_type", "基金指数"))
        basic_filter.append(("or", "sec_type", "债券指数"))
        basic_filter.append(("or", "sec_type", "期货指数"))    #
        #
        documents_symbol = self.Find(self.dbname, "sec_basic_info",
                                     filter=basic_filter,
                                     projection=["sec_issuer_id", "sec_code", "thscode", "sec_id"])
        df_symbol = Gadget.DocumentsToDataFrame(documents_symbol)

        documents_instrument = self.Find(self.dbname, "index_basic_info", filter=filter,
                                         projection=["index_code", "index_short_name", "index_name", "base_period", "enable_date", "index_id"])
        df_instrument = Gadget.DocumentsToDataFrame(documents_instrument)
        #
        df_instrument = pd.merge(df_instrument, df_symbol, how="left", left_on="index_code", right_on="sec_code")
        df_instrument["symbol"] = df_instrument["thscode"]
        return df_instrument

    def _Get_Instruments_DataFrame_Stock(self, filter=[], projection=[], listed=True):
        basic_filter = []
        basic_filter.append(("sec_type", "A股"))

        documents_symbol = self.Find(self.dbname, "sec_basic_info",
                                     filter=basic_filter,
                                     projection=["sec_issuer_id", "sec_code", "thscode", "sec_id", "sec_short_name_cn", "sec_full_name", "sec_type",\
                                                 "td_mkt", "listed_date", "issue_org_id", "listed_status"])
        df_symbol = Gadget.DocumentsToDataFrame(documents_symbol)

        df_symbol["symbol"] = df_symbol["thscode"]
        return df_symbol

    def _Get_Instruments_DataFrame_Bond(self, filter=[], projection=[], listed=True):
        bond_code = self.Find(self.dbname, "sec_classi_public_code_table", filter=[("class_encode", "002")], projection=["ctgry_code"])
        bond_code_res = [item[key] for item in bond_code for key in item]
        basic_filter = list(zip(['sec_type_code'] * len(bond_code_res), bond_code_res))
        basic_filter[0] = basic_filter[0][1:]

        documents_symbol = self.Find(self.dbname, "sec_basic_info",
                                     filter=basic_filter,
                                     projection=["sec_issuer_id", "sec_code", "thscode", "sec_id"])
        df_symbol = Gadget.DocumentsToDataFrame(documents_symbol)

        documents_instrument = self.Find(self.dbname, "bond_basic_info", filter=filter,
                                         projection=["bond_name", "interest_date", "maturity_date", "bond_id"])
        df_instrument = Gadget.DocumentsToDataFrame(documents_instrument)
        #
        df_instrument = pd.merge(df_instrument, df_symbol, how="left", left_on="sec_issuer_id", right_on="bond_id")
        df_instrument["symbol"] = df_instrument["thscode"]
        return df_instrument

    # def Get_Daily_Bar(self, symbol, instrument_type, projection=[], datetime1=None, datetime2=None):
    #     if datetime2:
    #         modified_datetime2 = datetime2 + datetime.timedelta(seconds=-1)
    #     df_bars = self.Get_Daily_Bar_DataFrame(symbol, instrument_type, projection=projection, datetime1=datetime1, datetime2=modified_datetime2)
    #     documents = df_bars.to_dict(orient='records')
    #     return documents

    def Get_Daily_Bar_DataFrame(self, symbol, instrument_type, projection=[], datetime1=None, datetime2=None):
        print("---THSDatabase::Get_Daily_Bar_DataFrame---", symbol, instrument_type, datetime1, datetime2)

        # 更新映射表
        self._Refresh_Instrument_Map_Table()

        # 转换小写
        if instrument_type:
            instrument_type = instrument_type.lower()
        #
        df_bars = pd.DataFrame()
        if instrument_type == "mutualfund":
            df_bars = self.Get_MutualFund_Dailybars_Dataframe(symbol, projection, datetime1, datetime2)
        elif instrument_type == "index":
            df_bars = self.Get_Index_Dailybars_Dataframe(symbol, projection, datetime1, datetime2)
        elif instrument_type == "stock":
            df_bars = self.Get_Stock_Dailybars_Dataframe(symbol, projection, datetime1, datetime2)
        elif instrument_type == "bond":
            df_bars = self.Get_Bond_Dailybars_Dataframe(symbol, projection, datetime1, datetime2)
        elif instrument_type == "future":
            pass

        #
        if len(df_bars) > 0:
            df_bars["symbol"] = symbol
            df_bars["datetime"] = pd.to_datetime(df_bars["date"])
        df_bars = self._Process_Projection(df_bars, projection)
        return df_bars

    def _Process_Projection(self, df, projection):
        real_projection = []
        if len(projection):
            for field in projection:
                if field in df.columns:
                    real_projection.append(field)
            return df[real_projection].copy(deep=True)
        else:
            return df

    def Get_MutualFund_Dailybars_Dataframe(self, symbol, out_projection, datetime1, datetime2):

        fund_code = self.fund_code_by_symbol.get(symbol)
        if fund_code == None:
            print("Can't Find thscode", fund_code, "of Symbol", symbol)
        #
        projection = ["ed", "unit_nv", "accum_nv", "adjust_nv", "fund_af"]
        filter = [("fund_code", fund_code)]
        #
        if datetime1:
            filter.append(("ed", ">=", datetime1))
        if datetime2:
            filter.append(("ed", "<=", datetime2))
        sort = [("ed", 1)]
        #
        dailybars = self.Find(self.dbname, "fund_net_value", filter=filter, projection=projection, sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars, rename={"ed": "date"})

        # 必要的类型转换
        df_bars["date"] = pd.to_datetime(df_bars["date"].dt.date)
        df_bars["unit_nv"] = pd.to_numeric(df_bars["unit_nv"])
        df_bars["accum_nv"] = pd.to_numeric(df_bars["accum_nv"])
        df_bars["adjust_nv"] = pd.to_numeric(df_bars["adjust_nv"])
        df_bars["fund_af"] = pd.to_numeric(df_bars["fund_af"])
        #
        df_bars.rename(columns={"adjust_nv": "adjusted_net_asset_value"}, inplace=True)
        return df_bars

    def Get_Index_Dailybars_Dataframe(self, symbol, out_projection, datetime1, datetime2):
        sec_id = self.sec_id_by_symbol.get(symbol)
        if sec_id == None:
            # print("Can't Find thscode of Symbol", symbol)
            return pd.DataFrame()
        #
        projection = ["td_date", "open_index", "high_index", "low_index", "latest_index", "vol", "trans_amt"]
        filter = [("sec_id", sec_id)]
        #
        if datetime1:
            filter.append(("td_date", ">=", datetime1))
        if datetime2:
            filter.append(("td_date", "<=", datetime2))
        sort = [("td_date", 1)]
        #
        dailybars = self.Find(self.dbname, "index_daily_quotation", filter=filter, projection=projection, sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars,
                                              rename={"td_date": "date",
                                                      "open_index": "open",
                                                      "low_index": "low",
                                                      "high_index": "high",
                                                      "latest_index": "close"})
        # 必要的类型转换
        df_bars["date"] = pd.to_datetime(df_bars["date"].dt.date)
        df_bars["close"] = pd.to_numeric(df_bars["close"])

        return df_bars

    def Get_Stock_Dailybars_Dataframe(self, symbol, out_projection, datetime1, datetime2):
        sec_id = self.sec_id_by_symbol.get(symbol)
        if sec_id is None:
            return pd.DataFrame()

        projection = ["td_date", "pre_close", "open_price", "high_trans", "low_trans", "latest_trans", "trans_vol", "trans_amt_rmb"]
        filter = [("sec_id", sec_id)]

        if datetime1:
            filter.append(("td_date", ">=", datetime1))
        if datetime2:
            filter.append(("td_date", "<=", datetime2))
        sort = [("td_date", 1)]
        #
        dailybars = self.Find(self.dbname, "shsz_stock_daily_quotation", filter=filter, projection=projection, sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars,
                                              rename={"td_date": "date",
                                                      "open_price": "open",
                                                      "high_trans": "high",
                                                      "low_trans": "low",
                                                      "latest_trans": "close",
                                                      "pre_close": "pre_close",
                                                      "trans_vol": "volume",
                                                      "trans_amt_rmb": "amount"})
        # 必要的类型转换
        df_bars["date"] = pd.to_datetime(df_bars["date"].dt.date)
        df_bars["close"] = pd.to_numeric(df_bars["close"])

        return df_bars

    def Get_Bond_Dailybars_Dataframe(self, symbol, out_projection, datetime1, datetime2):

        sec_id = self.sec_id_by_symbol.get(symbol)
        if sec_id is None:
            return pd.DataFrame()

        projection = ["td_date", "pre_close_full_price", "open_full_price", "high_full_price", "low_full_price", "clos_full_price", "vol", "amt"]
        filter = [("sec_id", sec_id)]

        if datetime1:
            filter.append(("td_date", ">=", datetime1))
        if datetime2:
            filter.append(("td_date", "<=", datetime2))
        sort = [("td_date", 1)]
        #
        dailybars = self.Find(self.dbname, "bond_daily_quotation", filter=filter, projection=projection, sort=sort)

        if len(dailybars) == 0:
            return pd.DataFrame()

        df_bars = Gadget.DocumentsToDataFrame(dailybars,
                                              rename={"td_date": "date",
                                                      "open_full_price": "open",
                                                      "high_full_price": "high",
                                                      "low_full_price": "low",
                                                      "clos_full_price": "close",
                                                      "pre_close_full_price": "pre_close",
                                                      "vol": "volume",
                                                      "amt": "amount"})
        # 必要的类型转换
        df_bars["date"] = pd.to_datetime(df_bars["date"].dt.date)
        df_bars["close"] = pd.to_numeric(df_bars["close"])
        return df_bars

def Test_THSDatabase(database):
    # intruments = database.Get_Instruments(instrument_type="MutualFund")
    # total_count = 10
    # count = 0
    # for instrument in intruments:
    #     symbol = instrument["symbol"]
    #     if symbol:
    #         df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="MutualFund")
    #         print(symbol)
    #         print(df_bars)
    #         count += 1
    #     if count >= total_count:
    #         break

    #
    intruments = database.Get_Instruments(instrument_type="Index")
    total_count = 10
    count = 0
    for instrument in intruments:
        symbol = instrument["symbol"]
        if symbol:
            df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="Index")
            print(symbol)
            print(df_bars)
            count += 1
        if count >= total_count:
            break


def Test_ETF(database):
    symbols = ["150008.SZ", "150009.SZ"]
    for symbol in symbols:
        df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="MutualFund")
        print(symbol)
        print(df_bars)


def Test_Stock(database):
    intruments = database.Get_Instruments(instrument_type="Stock")
    total_count = 10
    count = 0
    for instrument in intruments:
        symbol = instrument["symbol"]
        if symbol:
            df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="Stock")
            print(symbol)
            print(df_bars)
            count += 1
        if count >= total_count:
            break

def Test_Bond(database):
    intruments = database.Get_Instruments(instrument_type="Bond")
    total_count = 2
    count = 0
    for instrument in intruments:
        symbol = instrument["symbol"]
        if symbol:
            df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="Bond")
            print(symbol)
            print(df_bars)
            count += 1
        if count >= total_count:
            break

def Test_Bond2(database):
    # symbols = ["038006.IB", "122749.SH"]
    symbols = ["070207.IB","078012.IB","078011.IB","078010.IB","100706.SZ","010706.SH","070006.IB","070202.IB"]
    for symbol in symbols:
        df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="Bond")
        print(symbol)
        print(df_bars)


def Test_Index(database):
    symbols = ["000300.SH", "HSCI.HI", "000002.SH"]
    for symbol in symbols:
        df_bars = database.Get_Daily_Bar_DataFrame(symbol, instrument_type="Index")
        print(symbol)
        print(df_bars)


if __name__ == '__main__':
    #
    from Core.Config import *
    pathfilename = os.getcwd() + "\..\Config\config_develop.json"
    config = Config(pathfilename)
    database = config.DataBase("THS", config_field="TongHuaShunDB")
    realtime = config.RealTime()
    #
    # Test_THSDatabase(database)
    # Test_ETF(database)
    # Test_Stock(database)
    Test_Index(database)
    # Test_Bond2(database)